#!/usr/bin/env python
# encoding:utf-8
# @Time   : 2020/2/7
# @Author : 茶葫芦
# @Site   : 
# @File   : hotsong_net.py

import networkx.algorithms.community as community
import os
import sys
from prepare_graph_songs import GraphSongs
import argparse
import json
import networkx as nx
import matplotlib.pyplot as plt
from louvain import louvain_heuristices,modularity
from infomap import infomap_community

if __name__ == '__main__':

    jsonpath = os.path.dirname(os.path.abspath(__file__)) + "/../../data/data.json"
    json_str = json.load(open(jsonpath))
    gs_obj = GraphSongs(json_str)
    graph_songs = gs_obj.prepare_graph(list(json_str.keys()))
    graph_songs.add_nodes_from(['hotsong','coldsong'])
    for n in json_str:
        hotsong = json_str[n].get('hotSong',False)
        if hotsong and 'true' in hotsong:
            graph_songs.add_edge('hotsong', json_str[n]['songName'][0])
        else:
            graph_songs.add_edge('coldsong', json_str[n]['songName'][0])

    partition = louvain_heuristices(graph_songs)
    print('partition:', partition)
    print('modularity:', modularity(graph_songs, partition))
    # partition_infomap = infomap_community(graph_songs)
    # print('partition_infomap:', partition_infomap)
    # print('partition_infomap modularity:', modularity(graph_songs, partition_infomap))









